package com.mdnote.practice.sort;

import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;

/**
 * @author Rhythm-2019
 * @version 1.0
 * @date 2020/11/9 18:58
 * @Email rhythm_2019@163.com
 * @Description 973. 最接近原点的 K 个点
 */
public class LeetCode973 {

//    public int[][] kClosest(int[][] points, int K) {
//        if (points == null || points.length == 0 || K == 0) {
//            return new int[0][0];
//        }
//        // heap
//        PriorityQueue<Point> queue = new PriorityQueue<>();
//
//        for (int i = 0; i < points.length; i++) {
//            queue.offer(new Point(points[i][0], points[i][1]));
//        }
//
//        int[][] res = new int[K][2];
//        for (int i = 0; i < res.length; i++) {
//            Point point = queue.poll();
//            res[i][0] = point.x;
//            res[i][1] = point.y;
//        }
//        return res;
//    }
//
//    private static class Point implements Comparable<Point> {
//
//        private int x;
//        private int y;
//
//        public Point(int x, int y) {
//            this.x = x;
//            this.y = y;
//        }
//
//        @Override
//        public int compareTo(Point p) {
//            return (this.x * this.x + this.y * this.y) - (p.x * p.x + p.y * p.y);
//        }
//    }
//
    public int[][] kClosest(int[][] points, int K) {
        if (points == null || points.length == 0 || K == 0) {
            return new int[0][0];
        }
        // Quick Sort
        quickSort(points, 0, points.length - 1, K);

        return Arrays.copyOf(points, K);

    }

    private void quickSort(int[][] points, int start, int end, int k) {

        if (start >= end) {
            return;
        }

        int pivot = partition(points, start, end);
        // 当前pivot左边的元素到原点距离全部小于pivot，右边反之
        if (k < pivot - start + 1) {
            quickSort(points, start, pivot - 1, k);
        }
        else if (k > pivot - start + 1) {
            quickSort(points, pivot + 1, end, k - (pivot - start + 1) );
        }
    }

    private int partition(int[][] points, int start, int end) {
        // 默写模板
        int pivot = end, counter = start;
        for (int i = start; i < end; i++) {
            if (compareLen(points, i, pivot)) {
                swap(points, i, counter);
                counter++;
            }
        }
        swap(points, pivot, counter);
        return counter;
    }

    private void swap(int[][] points, int p1, int p2) {
        int[] tmp = points[p1];
        points[p1] = points[p2];
        points[p2] = tmp;
    }

    private boolean compareLen(int[][] points, int p1, int p2) {
        // p1 是否比 o2 近
        return points[p1][0] * points[p1][0] + points[p1][1] * points[p1][1] <= points[p2][0] * points[p2][0] + points[p2][1] * points[p2][1];
    }

    @Test
    public void test_01(){
        LeetCode973 leetCode973 = new LeetCode973();

        int[][] points = {
                {1, 3},
                {2, -2},
                {-2, 2},
        };
//        Assert.assertArrayEquals(new int[][]{{1, 1}}, leetCode973.kClosest(points, 1));
        int[][] actuals = leetCode973.kClosest(points, 2);
        System.out.println(Arrays.toString(actuals[0]));
        System.out.println(Arrays.toString(actuals[1]));
        Assert.assertArrayEquals(new int[][]{{-2, 2}, {2, -2}}, actuals);
    }
}
